#include "hmm.h"

#include <QFile>

#include <iostream>
using namespace std;

using namespace mlgesture;

Hmm::Hmm() {
	statesCount_ = -1;
	observationCount_ = -1;
}

Hmm::Hmm(int states, int observation, ModelConnection mode) {
	setParams(states, observation, mode);
}

Hmm::~Hmm() {}

void Hmm::setParams(int states, int observation, ModelConnection mode) throw(HandGestureException) {
	statesCount_ = states;
	observationCount_ = observation;
	observation_.resize(states);
	for (int i = 0; i<states; i++) {
		observation_[i].resize(observation);
		for (int j=0; j<observation; j++) {
			observation_[i][j] = 1.0 / observation;
		}
	}
	transition_.resize(states);
	for (int i = 0; i<states; i++) {
		transition_[i].resize(states);
		for (int j=0; j<states; j++) {
			transition_[i][j] = 0.0;
		}
	}
	if (mode == LINEAR_MODEL) {
		double prob;
		for (int i=0; i<states; i++) {
			if (i < states - 2) {
				prob = 1.0/3;
				transition_[i][i] = prob;
				transition_[i][i+1] = prob;
				transition_[i][i+2] = prob;
			}
			else if (i == states - 2) {
				prob = 1.0 / 2;
				transition_[i][i] = prob;
				transition_[i][i+1] = prob;
			}
			else {
				transition_[i][i] = 1.0;
			}
		}
	}
	else if (mode == FULL_CONNECTED_MODEL) {
		double prob = 1. / states;
		for(int i=0; i<states; i++) {
			for(int j=0; j<states; j++) {
				transition_[i][j] = prob;
			}
		}
	}
	else {
		throw HandGestureException("Unrecognized HMM mode");
	}
	pi_.resize(states);
	pi_[0] = 1.0;
	for(int i=1; i<states; i++) {
		pi_[i] = 0.0;
	}
}

double Hmm::viterbi(const vector<int> &data) const {
	vector<double> G(statesCount_), Glast(statesCount_);
	for(int i=0; i<statesCount_; i++) {
		Glast[i] = pi_[i];
	}
	for (size_t i=0; i<data.size(); i++) {
		for (int j=0; j<statesCount_; j++) {
			double max = -1;
			for (int k=0; k<statesCount_; k++) {
				if (Glast[k] * transition_[k][j]>max) {
					max = Glast[k] * transition_[k][j];
				}
			}
			G[j] = max * observation_[j][data[i]];
		}
		Glast = G;
	}
	double max = -1;
	for (int i=0; i < statesCount_; i++) {
		G[i] = Glast[i];
		if (G[i] > max) {
			max = G[i];
		}
	}
	return max;
}

void Hmm::baumWelsh(const vector<vector<int> > &data) throw(HandGestureException) {
	if (data.size() == 0) {
		throw HandGestureException("Training data is empty");
	}
	double last_prob = 0;

	vector<vector<double> > new_transition(statesCount_);
	vector<vector<double> > new_observation(statesCount_);
	vector<double> new_pi(statesCount_);
	for (int i=0; i<statesCount_; i++) {
		new_transition[i].resize(statesCount_);
		new_observation[i].resize(observationCount_);
	}
	vector<vector<double> > last_transition;
	vector<vector<double> > last_observation;
	vector<double> last_pi;

	int iteration = 0;
	while (1) {
		iteration++;
		for (int i=0; i<statesCount_; i++) {
			new_pi[i] = 0;
			for (int j=0; j<statesCount_; j++) {
				new_transition[i][j] = 0;
			}
			for (int j=0; j<observationCount_; j++) {
				new_observation[i][j] = 0;
			}
		}

		for (size_t z=0; z<data.size(); z++) {
			vector<vector<double> > A, B;
			A.resize(data[z].size());
			for (size_t i=0; i<data[z].size(); i++) {
				A[i].resize(statesCount_);
			}
			for (int i=0; i<statesCount_; i++) {
				A[0][i] = pi_[i]*observation_[i][data[z][0]];
				if (A[0][i] < 0) A[0][i] = 0;
			}
			for (size_t i=1; i<data[z].size(); i++) {
				for (int j=0; j<statesCount_; j++) {
					A[i][j] = 0;
					for (int k=0; k<statesCount_; k++) {
						A[i][j] += A[i-1][k]*transition_[k][j];
					}
					A[i][j] *= observation_[j][data[z][i]];
					if (A[i][j] < 0) A[i][j] = 0;
				}
			}
			B.resize(data[z].size()+1);
			for (size_t i=0; i<data[z].size()+1; i++) {
				B[i].resize(statesCount_);
			}
			for (int i=0; i<statesCount_; i++) {
				B[data[z].size()][i] = 1;
				B[data[z].size() - 1][i] = 1;
			}
			for (int i= (int)data[z].size() - 2; i>=0; i--) {
				for (int j=0; j<statesCount_; j++) {
					B[i][j] = 0;
					for (int k=0; k<statesCount_; k++) {
						B[i][j] += transition_[j][k]*observation_[k][data[z][i+1]]*B[i+1][k];
					}
					if (B[i][j] < 0) B[i][j] = 0;
				}
			}
			vector<vector<vector<double> > > E(data[z].size());
			for (size_t i=0; i<data[z].size(); i++) {
				double x = 0;
				for (int j=0; j<statesCount_; j++) {
					x += A[i][j]*B[i][j];
				}
				E[i].resize(statesCount_);
				for (int j=0; j<statesCount_; j++) {
					E[i][j].resize(statesCount_);
					for (int k=0; k<statesCount_; k++) {
						E[i][j][k] = A[i][j]*transition_[j][k]*B[i+1][k];
						if (i!=data[z].size() - 1) {
							E[i][j][k] *= observation_[k][data[z][i+1]];
						}
						E[i][j][k] /= x;
						if (E[i][j][k]<0) E[i][j][k] = 0;
					}
				}
			}
			for (int i=0; i<statesCount_; i++) {
				for (int j=0; j<statesCount_; j++) {
					new_pi[i] += E[0][i][j];
				}
			}

			for (int i=0; i<statesCount_; i++) {
				double x = 0;
				for (int j=0; j<statesCount_; j++) {
					for (size_t k=0; k<data[z].size(); k++) {
						x += E[k][i][j];
					}
				}
				for (int j=0; j<statesCount_; j++) {
					double y = 0;
					for (size_t k=0; k<data[z].size(); k++) {
						y += E[k][i][j];
					}
					if (y/x > 0)
						new_transition[i][j] += y/x;
				}
			}
			for (int i=0; i<statesCount_; i++) {
				double x = 0;
				for (int j=0; j<statesCount_; j++) {
					for (size_t k=0; k<data[z].size(); k++) {
						x += E[k][i][j];
					}
				}
				for (int j=0; j<observationCount_; j++) {
					double y = 0;
					for (size_t k=0; k<data[z].size(); k++) {
						if (data[z][k] == j) {
							for (int l = 0; l<statesCount_; l++) {
								y += E[k][i][l];
							}
						}
					}
					if (y/x>0)
						new_observation[i][j] += y/x;
				}
			}
		}
		last_transition = transition_;
		last_observation = observation_;
		last_pi = pi_;

		for (int i=0; i<statesCount_; i++) {
			for (int j=0; j<statesCount_; j++) {
				transition_[i][j] = new_transition[i][j]/data.size();
			}
		}
		for (int i=0; i<statesCount_; i++) {
			for (int j=0; j<observationCount_; j++) {
				observation_[i][j] = new_observation[i][j]/data.size();
			}
		}
		for (int i=0; i<statesCount_; i++) {
			pi_[i] = new_pi[i]/data.size();
		}

		double prob_sum = 0;
		for (size_t i = 0; i<data.size(); i++) {
			double prob = viterbi(data[i]);
			prob_sum += prob;
		}
		if (prob_sum / data.size() - last_prob < 0) {
			transition_ = last_transition;
			observation_ = last_observation;
			pi_ = last_pi;
			break;
		}
		last_prob = prob_sum / data.size();
		if (iteration > 5000) break;
	}
	double prob_sum = 0;
	for (size_t i = 0; i<data.size(); i++) {
		double prob = viterbi(data[i]);
		prob_sum += prob;
	}
}

bool Hmm::read(QString file) throw(HandGestureException) {
	bool ok1, ok2;
	bool is_state = false;
	int act_state = 0;
	int second_state, observation_id;
	double prob;
	QXmlStreamReader::TokenType token;
	QFile qfile(file);
	if (!qfile.open(QFile::ReadOnly | QFile::Text))
		throw HandGestureException("Cannot open " + file.toStdString());
	xmlReader_.setDevice(&qfile);
	if (xmlReader_.readNextStartElement()) {
		observationCount_ = xmlReader_.attributes().value("observations").toString().toInt(&ok1);
		statesCount_ = xmlReader_.attributes().value("states").toString().toInt(&ok2);
		if ((xmlReader_.name() != "hmm") || !ok1 || !ok2) {
			xmlReader_.raiseError(QObject::tr("The file is not an xml-hmm file."));
			return false;
		}

		setParams(statesCount_, observationCount_);

		while ((token = xmlReader_.readNext()) != QXmlStreamReader::EndDocument) {
			if (token == QXmlStreamReader::StartElement) {
				if (xmlReader_.name() == "state") {
					if (is_state) {
						xmlReader_.raiseError("State in state token");
						return false;
					}
					act_state = xmlReader_.attributes().value("id").toString().toInt(&ok1);
					if (!ok1) {
						xmlReader_.raiseError("No state id");
						return false;
					}
					is_state = true;
				}
				else if (xmlReader_.name() == "transition") {
					if (!is_state) {
						xmlReader_.raiseError("Transition outside state");
						return false;
					}
					second_state = xmlReader_.attributes().value("id").toString().toInt(&ok1);
					if (!ok1) {
						xmlReader_.raiseError("No transition state id");
						return false;
					}
					prob = xmlReader_.readElementText().toDouble(&ok1);
					if (!ok1) {
						xmlReader_.raiseError("No transition propability");
						return false;
					}
					transition_[act_state][second_state] = prob;
				}
				else if (xmlReader_.name() == "observation") {
					if (!is_state) {
						xmlReader_.raiseError("Observation outside state");
						return false;
					}
					observation_id = xmlReader_.attributes().value("id").toString().toInt(&ok1);
					if (!ok1) {
						xmlReader_.raiseError("No observation state id");
						return false;
					}
					prob = xmlReader_.readElementText().toDouble(&ok1);
					if (!ok1) {
						xmlReader_.raiseError("No observation propability");
						return false;
					}
					observation_[act_state][observation_id] = prob;
				}
				else if (xmlReader_.name() == "pi") {
					if (!is_state) {
						xmlReader_.raiseError("Pi outside state");
						return false;
					}
					prob = xmlReader_.readElementText().toDouble(&ok1);
					if (!ok1) {
						xmlReader_.raiseError("No pi propability");
						return false;
					}
					pi_[act_state] = prob;
				}
				else xmlReader_.skipCurrentElement();
			}
			else if (token == QXmlStreamReader::EndElement) {
				if (xmlReader_.name() == "state") {
					is_state = false;
				}
			}
		}
	}
	return !xmlReader_.error();
}

void Hmm::write(QString file) throw(HandGestureException) {
	QFile qfile(file);
	if (!qfile.open(QFile::WriteOnly | QFile::Text))
		throw HandGestureException("Cannot open " + file.toStdString());
	xmlWriter_.setAutoFormatting(true);
	xmlWriter_.setDevice(&qfile);
	xmlWriter_.writeStartDocument();
	xmlWriter_.writeStartElement("hmm");
	xmlWriter_.writeAttribute("states",QString::number(statesCount_));
	xmlWriter_.writeAttribute("observations", QString::number(observationCount_));

	for (int i=0; i<statesCount_; i++) {
		xmlWriter_.writeStartElement("state");
		xmlWriter_.writeAttribute("id", QString::number(i));
		for (int j=0; j<statesCount_; j++) {
			xmlWriter_.writeStartElement("transition");
			xmlWriter_.writeAttribute("id", QString::number(j));
			xmlWriter_.writeCharacters(QString::number(transition_[i][j]));
			xmlWriter_.writeEndElement();
		}
		for (int j=0; j<observationCount_; j++) {
			xmlWriter_.writeStartElement("observation");
			xmlWriter_.writeAttribute("id", QString::number(j));
			xmlWriter_.writeCharacters(QString::number(observation_[i][j]));
			xmlWriter_.writeEndElement();
		}
		xmlWriter_.writeTextElement("pi", QString::number(pi_[i]));
		xmlWriter_.writeEndElement();
	}

	xmlWriter_.writeEndDocument();
}

Hmm* Hmm::makeGarbage(const vector<Hmm *> models) throw(HandGestureException) {
	int states_num = 0;
	int observation_num = -1;
	for (size_t i=0; i<models.size(); i++) {
		states_num += models[i]->statesCount_;
		if (observation_num == -1) {
			observation_num = models[i]->observationCount_;
		}
		else {
			if (observation_num != models[i]->observationCount_)
				throw HandGestureException("Different observation number in models");
		}
	}

	Hmm* model = new Hmm(states_num, observation_num, FULL_CONNECTED_MODEL);

	int global_state = 0;
	int begin_state = 0;
	for (size_t i = 0; i < models.size(); i++) {
		for (int state = 0; state<models[i]->statesCount_; state++, global_state++) {
			for (int j=0; j<states_num; j++) {
				model->transition_[global_state][j] = 1.0 / (states_num - 1);
			}
			for (int j=0; j<models[i]->statesCount_; j++) {
				model->transition_[global_state][begin_state+j] = (1.0 - models[i]->transition_[state][j])/(states_num - 1);
			}
			for (int j=0; j<observation_num; j++) {
				model->observation_[global_state][j] = models[i]->observation_[state][j];
			}
			model->pi_[global_state] = 1.0 / states_num;
		}
		begin_state += models[i]->statesCount_;
	}

	return model;
}


int Hmm::getStatesNum() const {
	return statesCount_;
}
